import numpy as np
from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection

# 连接到 Milvus
connections.connect("default", host="localhost", port="19530")

# 定义字段和 Schema
field_dim = 128
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=field_dim)
]
schema = CollectionSchema(fields, "description")

# 集合名称
collection_name = "example_collection"

# 检查集合是否存在
if utility.has_collection(collection_name):
    # 如果存在，直接加载集合
    collection = Collection(name=collection_name)
    print(f"Collection {collection_name} already exists. Using existing collection.")
    collection.load()  # 加载集合到内存

else:
    # 如果不存在，创建新集合
    collection = Collection(name=collection_name, schema=schema)
    print(f"Collection {collection_name} created.")
    # 创建索引（如果需要 L2 距离）
    index_params = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128}
    }
    collection.create_index("vector", index_params)
    collection.load()  # 加载集合到内存


# 生成随机向量数据
num_vectors = 10
vectors = np.random.random((num_vectors, field_dim)).astype(np.float32).tolist()

# 插入数据，将序号作为 ID
data = [
    {"id": i, "vector": vectors[i]}  # 使用序号 i 作为 ID
    for i in range(num_vectors)
]
insert_result = collection.insert(data)
print(f"Inserted {len(data)} vectors with IDs: {insert_result.primary_keys}")

# 刷新集合
collection.flush()

# 关闭连接
#connections.disconnect("default")

# 生成查询向量（随机选择一个向量作为查询）
query_vector = np.random.random((1, field_dim)).astype(np.float32).tolist()

# 执行相似度搜索
search_params = {"metric_type": "L2", "params": {"nprobe": 10}}  # 使用 L2 距离
results = collection.search(
    data=query_vector,
    anns_field="vector",
    param=search_params,
    limit=3,  # 返回最相似的 3 个结果
    output_fields=["id"],  # 返回的字段（这里是 id）
    expr=None  # 可选的过滤表达式
)

# 打印搜索结果
for hits in results:
    for hit in hits:
        print(f"ID: {hit.entity.get('id')}, Distance: {hit.distance}")

# 关闭连接
connections.disconnect("default")